Description

This PR fixes a compatibility issue with recent vLLM changes that now require model classes to implement a get_input_embeddings() method.
Without this method, vLLM fails its interface validation during model registration, breaking TPU model integration.

To address this, we add a dummy get_input_embeddings() implementation to the vLLM-compatible wrapper class in
tpu_inference/models/common/model_loader.py.
Similar to the existing dummy forward() method, this implementation only satisfies vLLM’s type checks and raises
NotImplementedError if invoked. This prevents JAX model initialization during import or introspection.

Why this change is needed

• vLLM recently introduced a strict requirement for model classes to define get_input_embeddings()
  (link).
• TPU inference uses a dummy PyTorch wrapper to register JAX models into vLLM’s registry.
• Since this wrapper lacked get_input_embeddings, vLLM failed model registration checks.

Implementation details

• Added unimplemented_get_input_embeddings() dummy function to the wrapper type.
• Registered it inside the dynamically created wrapper class.
• Added a test tests/test_vllm_wrapper.py to ensure:
  • The wrapper defines get_input_embeddings().
  • The method raises NotImplementedError.
  • The class passes is_vllm_model() validation.

Related Issue

Fixes: #951

Tests

pytest -v tests/models/common/test_model_loader.py 

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.